import os
import logging
import pickle
import re
from PyPDF2 import PdfReader
from langchain.chains.question_answering import load_qa_chain
from langchain_openai import ChatOpenAI
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.callbacks.manager import get_openai_callback
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from typing import List, Tuple




def extract_text_with_page_numbers(pdf) -> Tuple[str, List[int]]:
    """
    从PDF中提取文本并记录每行文本对应的页码

    参数:
        pdf: PDF文件对象

    返回:
        text: 提取的文本内容
        page_numbers: 每行文本对应的页码列表
    """
    # 初始化空字符串用于存储提取的所有文本内容
    text = ""
    # 初始化空列表用于存储每行文本对应的页码
    page_numbers = []

    # 遍历PDF中的每一页，enumerate从1开始编号（页码从1开始）
    for page_number, page in enumerate(pdf.pages, start=1):
        # 从当前页面提取文本内容
        extracted_text = page.extract_text()
        # 检查提取的文本是否为空
        if extracted_text:
            # 将提取的文本追加到总文本中
            text += extracted_text
            # 计算当前页的行数并将其页码添加到page_numbers列表中
            # split("\n")将文本按换行符分割，得到行数
            page_numbers.extend([page_number] * len(extracted_text.split("\n")))
        else:
            # 如果当前页面没有提取到文本，记录警告日志
            logging.warning(f"No text found on page {page_number}.")

    # 使用正则表达式去除"百度文库  - 好好学习，天天向上"文本
    # 这个文本在文档中出现，但对实际内容没有意义，需要清除
    text = re.sub(r"百度文库\s*-\s*好好学习，天天向上\s*", "", text)

    # 返回提取的文本内容和对应的页码列表
    return text, page_numbers




def split_text_into_chunks(text: str, page_numbers: List[int]) -> Tuple[List[str], dict]:
    """
    将文本分割成小块并创建页码映射

    参数:
        text: 要分割的文本
        page_numbers: 每行文本对应的页码列表

    返回:
        chunks: 分割后的文本块列表
        page_info: 文本块到页码的映射字典
    """
    # 创建文本分割器，用于将长文本分割成小块
    # separators: 定义分割符的优先级列表，按顺序尝试分割
    # chunk_size: 每个文本块的最大字符数
    # chunk_overlap: 相邻文本块之间的重叠字符数，避免内容截断
    # length_function: 用于计算文本长度的函数
    text_splitter = RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", ".", " ", ""],
        chunk_size=512,
        chunk_overlap=128,
        length_function=len,
    )

    # 使用文本分割器将输入文本分割成多个较小的文本块
    chunks = text_splitter.split_text(text)
    # logging.debug(f"Text split into {len(chunks)} chunks.")
    print(f"文本被分割成 {len(chunks)} 个块。")

    # 创建页码信息字典，将每个文本块映射到其原始页码
    # 使用enumerate遍历chunks列表，将每个文本块与其在列表中的索引关联
    # 通过索引从page_numbers列表中获取对应的页码
    page_info = {chunk: page_numbers[i] for i, chunk in enumerate(chunks)}

    return chunks, page_info







def create_embeddings_and_vectorstore(chunks: List[str]) -> FAISS:
    """
    将文本块转换为向量并创建向量存储

    参数:
        chunks: 文本块列表

    返回:
        knowledgeBase: FAISS向量存储对象
    """
    # 创建嵌入模型实例，使用阿里百炼平台的文本嵌入模型
    # model: 指定使用的嵌入模型名称
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v2"
    )
    # 使用分割后的文本块和嵌入模型创建FAISS向量数据库
    # FAISS是一种高效的相似性搜索库，用于快速检索相似的文本向量
    knowledgeBase = FAISS.from_texts(chunks, embeddings)
    print("已从文本块创建知识库...")

    return knowledgeBase






def save_vectorstore(knowledgeBase: FAISS, page_info: dict, save_path: str):
    """
    保存向量存储和页码信息到磁盘

    参数:
        knowledgeBase: FAISS向量存储对象
        page_info: 页码信息字典
        save_path: 保存路径
    """
    # 如果提供了保存路径，则将向量数据库和页码信息持久化存储
    if save_path:
        # 确保指定的保存目录存在，如果不存在则创建，exist_ok=True避免已存在时报错
        os.makedirs(save_path, exist_ok=True)

        # 将FAISS向量数据库保存到指定路径
        # save_local是FAISS提供的方法，用于持久化存储向量数据库
        knowledgeBase.save_local(save_path)
        print(f"向量数据库已保存到: {save_path}")

        # 将页码信息序列化保存到同一目录下的page_info.pkl文件中
        # 使用pickle.dump将page_info字典序列化并写入文件
        with open(os.path.join(save_path, "page_info.pkl"), "wb") as f:
            pickle.dump(page_info, f)
        print(f"页码信息已保存到: {os.path.join(save_path, 'page_info.pkl')}")






def load_knowledge_base(load_path: str, embeddings=None) -> FAISS:
    """
    从磁盘加载向量数据库和页码信息

    参数:
        load_path: 向量数据库的保存路径
        embeddings: 可选，嵌入模型。如果为None，将创建一个新的DashScopeEmbeddings实例

    返回:
        knowledgeBase: 加载的FAISS向量数据库对象
    """
    # 如果没有提供嵌入模型，则创建一个新的DashScopeEmbeddings实例
    # 这确保即使在加载时也能正确初始化嵌入模型
    if embeddings is None:
        # 使用与创建时相同的text-embedding-v2模型创建嵌入实例
        embeddings = DashScopeEmbeddings(
            model="text-embedding-v2"
        )

    # 从指定路径加载FAISS向量数据库
    # allow_dangerous_deserialization=True允许反序列化，因为FAISS保存的数据需要此参数
    knowledgeBase = FAISS.load_local(load_path, embeddings, allow_dangerous_deserialization=True)
    print(f"向量数据库已从 {load_path} 加载。")

    # 构造页码信息文件的完整路径
    page_info_path = os.path.join(load_path, "page_info.pkl")
    # 检查页码信息文件是否存在
    if os.path.exists(page_info_path):
        # 以二进制读取模式打开页码信息文件
        with open(page_info_path, "rb") as f:
            # 使用pickle.load反序列化页码信息字典
            page_info = pickle.load(f)
        # 将加载的页码信息附加到知识库对象上
        knowledgeBase.page_info = page_info

    # 返回加载完成的知识库对象
    return knowledgeBase





def retrieve_relevant_documents(knowledgeBase: FAISS, query: str) -> List:
    """
    根据问题检索相关文档

    参数:
        knowledgeBase: FAISS向量存储对象
        query: 查询问题

    返回:
        docs: 相关文档列表
    """
    # 在知识库中执行相似度搜索，找到与查询最相关的文档块
    # similarity_search会返回与查询语义最相似的文档列表
    docs = knowledgeBase.similarity_search(query)
    return docs





def generate_answer_with_llm(docs: List, query: str) -> Tuple[str, str]:
    """
    使用大模型生成答案

    参数:
        docs: 相关文档列表
        query: 查询问题

    返回:
        response_text: 模型的回答文本
        cost_info: 成本信息
    """
    # 初始化对话大模型，使用阿里云百炼平台的Qwen模型
    chatLLM = ChatOpenAI(
        # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
        # 从环境变量中获取DASHSCOPE_API_KEY作为API密钥
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        # 设置API的基础URL，使用阿里云百炼平台的兼容模式
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        # 使用qwen-plus模型，这是通义千问的一个版本
        model="qwen-plus"
    )

    # 加载问答链，使用stuff类型的链将所有相关文档一次性传递给模型
    chain = load_qa_chain(chatLLM, chain_type="stuff")

    # 准备输入数据，包含查询问题和相关文档
    input_data = {"input_documents": docs, "question": query}

    # 使用回调函数跟踪API调用成本，get_openai_callback可以统计token使用量和费用
    with get_openai_callback() as cost:
        # 执行问答链，获取模型的回答
        response = chain.invoke(input=input_data)
        # 获取查询处理的成本信息
        cost_info = f"查询已处理。成本: {cost}"
        # 获取模型的输出文本回答
        response_text = response["output_text"]

    return response_text, cost_info




def main():
    """
    主函数：测试每个功能函数
    """
    # 创建全局的嵌入模型实例，确保创建和加载时使用相同的模型
    embeddings = DashScopeEmbeddings(model="text-embedding-v2")

    print("=== 测试PDF文本提取 ===")
    try:
        # 读取PDF文件，使用PyPDF2库创建PdfReader对象
        pdf_reader = PdfReader('./浦发上海浦东发展银行西安分行个金客户经理考核办法.pdf')
        # 调用extract_text_with_page_numbers函数提取PDF文本及对应页码
        text, page_numbers = extract_text_with_page_numbers(pdf_reader)
        print(f"成功提取文本，文本长度: {len(text)} 字符")
        print(f"页码数量: {len(page_numbers)}")
        print(f"前5个页码: {page_numbers[:5] if len(page_numbers) > 5 else page_numbers}")
    except Exception as e:
        print(f"PDF文本提取失败: {e}")
        return

    print("\n=== 测试文本切割 ===")
    try:
        # 文本切割
        chunks, page_info = split_text_into_chunks(text, page_numbers)
        print(f"文本被分割成 {len(chunks)} 个块")
        print(f"第一个文本块预览: {chunks[0][:100] if chunks else '无内容'}...")
    except Exception as e:
        print(f"文本切割失败: {e}")
        return

    print("\n=== 测试文本转向量 ===")
    try:
        # 文本转向量并创建向量存储
        knowledgeBase = create_embeddings_and_vectorstore(chunks)
        # 将页码信息附加到知识库对象上，便于后续检索时追踪来源页码
        knowledgeBase.page_info = page_info
        print("成功创建向量存储")
    except Exception as e:
        print(f"文本转向量失败: {e}")
        return

    print("\n=== 测试向量存储保存 ===")
    try:
        # 保存向量存储（可选）
        save_vectorstore(knowledgeBase, page_info, "./vector_db")
        print("向量存储保存成功")
    except Exception as e:
        print(f"向量存储保存失败: {e}")

    print("\n=== 测试向量存储加载 ===")
    try:
        # 加载向量存储，使用相同的嵌入模型
        loaded_knowledgeBase = load_knowledge_base("./vector_db", embeddings)
        print("向量存储加载成功")
    except Exception as e:
        print(f"向量存储加载失败: {e}")
        # 如果加载失败，使用之前创建的知识库
        loaded_knowledgeBase = knowledgeBase

    print("\n=== 测试文档检索 ===")
    try:
        # 设置查询问题
        query = "客户经理每年评聘申报时间是怎样的？"
        # 根据问题检索相关文档
        docs = retrieve_relevant_documents(loaded_knowledgeBase, query)
        print(f"检索到 {len(docs)} 个相关文档")
        if docs:
            print(f"第一个文档预览: {docs[0].page_content[:100] if docs else '无内容'}...")
        else:
            print("未检索到相关文档")
    except Exception as e:
        print(f"文档检索失败: {e}")
        import traceback
        traceback.print_exc()
        return

    print("\n=== 测试大模型生成答案 ===")
    try:
        # 使用大模型生成答案
        response_text, cost_info = generate_answer_with_llm(docs, query)
        print(cost_info)
        print(f"生成的答案预览: {response_text[:200] if response_text else '无内容'}...")
    except Exception as e:
        print(f"大模型生成答案失败: {e}")
        import traceback
        traceback.print_exc()
        return

    print("\n=== 所有函数测试完成 ===")


if __name__ == '__main__':
    main()